import cv2
import numpy as np
from scipy.signal.windows import gaussian


class EyeFundusProcessor:
    def __init__(self):
        pass

    def apply_clahe(self, image, clip_limit=2.0, tile_grid_size=(8, 8)):
        """
        对彩色图像的每个通道分别应用 CLAHE（对比度受限的自适应直方图均衡化）。
        """
        clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
        channels = cv2.split(image)
        clahe_channels = [clahe.apply(ch) for ch in channels]
        return cv2.merge(clahe_channels)

    def pre_process_channels(self, image):
        """
        将图像拆分成三个通道，并对红色（视盘）和绿色通道（视杯）进行预处理。
        """
        # 注意：图像输入为BGR格式，Abo为蓝色通道，Ago为绿色通道，Aro为红色通道
        Abo, Ago, Aro = cv2.split(image)

        # 预处理红色通道（视盘）
        Ar = Aro.astype(np.float32)
        Ar = Ar - Ar.mean()
        Ar = Ar - Ar.mean() - Ar.std()
        Ar = Ar - Ar.mean() - Ar.std()

        # 预处理绿色通道（视杯）
        Ag = Ago.astype(np.float32)
        Ag = Ag - Ag.mean()
        Ag = Ag - Ag.mean() - Ago.std()

        return Ar, Ag

    def compute_thresholds(self, Ar, Ag):
        """
        根据预处理后的通道数据计算阈值：
          - Thr 用于分割视盘（红色通道）
          - Thg 用于分割视杯（绿色通道）
        """
        Thr = Ar.std()  # 视盘阈值
        Thg = Ag.mean() + 2 * Ag.std() + 49.5 + 12  # 视杯阈值（原代码公式）
        return Thr, Thg

    def smooth_histogram(self, channel, gauss_filter):
        """
        对输入通道数据计算直方图并使用高斯窗进行平滑。
        """
        hist, bins = np.histogram(channel.ravel(), 256, [0, 256])
        smooth_hist = np.convolve(gauss_filter, hist, mode='same')
        return hist, smooth_hist

    def apply_threshold(self, channel, threshold, binary_value=255):
        """
        对单通道图像应用阈值处理，生成二值化掩码。
        """
        binary = np.zeros(channel.shape, dtype=np.uint8)
        binary[channel > threshold] = binary_value
        return binary

    def overlay_mask(self, original, mask, color=(0, 255, 0), alpha=0.4, invert=False):
        """
        将二值掩码以指定颜色和透明度叠加在原图上。

        参数 invert 控制：
          - invert=False（默认）：掩码中非零部分叠加透明绿色
          - invert=True：掩码中非零部分保持原图显示，其他部分叠加透明绿色
        """
        overlay = original.copy()
        mask_bool = mask.astype(bool)

        if invert:
            # 对于检测区域（mask==True），保留原图；其他区域叠加透明绿色
            inv_mask = ~mask_bool
            overlay[inv_mask] = (overlay[inv_mask] * (1 - alpha) + np.array(color) * alpha).astype(np.uint8)
        else:
            # 掩码中非零部分叠加透明绿色
            overlay[mask_bool] = (overlay[mask_bool] * (1 - alpha) + np.array(color) * alpha).astype(np.uint8)

        return overlay

    def draw_bounding_boxes(self, original, mask, box_color=(0, 255, 0), thickness=20, min_area=500):
        """
        在原图上对二值掩码中连通区域绘制边框。

        参数：
          - original: 原图（BGR格式）
          - mask: 二值化掩码（非零区域即为感兴趣区域）
          - box_color: 边框颜色（默认绿色）
          - thickness: 边框线宽
          - min_area: 连通区域最小面积（过滤噪声）
        """
        output = original.copy()
        # 查找掩码中的轮廓
        contours, _ = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in contours:
            if cv2.contourArea(cnt) >= min_area:
                x, y, w, h = cv2.boundingRect(cnt)
                cv2.rectangle(output, (x, y), (x + w, y + h), box_color, thickness)
        return output

    def process_image_from_array(self, image):
        """
        直接对输入的图像数组进行处理，完成 CLAHE、预处理、阈值分割以及直方图平滑，
        并返回原图、视盘掩码、视杯掩码以及直方图数据。
        """
        if image is None:
            raise ValueError("图像数据为空！")

        # 先对图像进行 CLAHE
        image_clahe = self.apply_clahe(image)

        # 后续操作基于 CLAHE 处理后的图像
        Ar, Ag = self.pre_process_channels(image_clahe)
        Thr, Thg = self.compute_thresholds(Ar, Ag)

        # 构造高斯滤波窗口（用于直方图平滑）
        gauss_filter = gaussian(99, std=6)
        gauss_filter = gauss_filter / np.sum(gauss_filter)

        # 计算平滑直方图（仅用于展示）
        hist_g, smooth_hist_g = self.smooth_histogram(Ag, gauss_filter)
        hist_r, smooth_hist_r = self.smooth_histogram(Ar, gauss_filter)

        # 应用阈值生成二值化掩码
        disc_mask = self.apply_threshold(Ar, Thr, binary_value=255)
        cup_mask = self.apply_threshold(Ag, Thg, binary_value=1)
        # 将 cup_mask 转换为 0/255 格式，方便后续处理显示
        cup_mask = (cup_mask * 255).astype(np.uint8)

        return image_clahe, disc_mask, cup_mask, (hist_g, smooth_hist_g), (hist_r, smooth_hist_r)

    def display_results(self, original, disc_mask, cup_mask, output_path=None, show_overlay=True, show_bbox=False,
                        alpha=0.4):
        """
        展示或保存处理结果，提供多种显示选项。

        参数：
          - original: 原始图像（BGR格式）
          - disc_mask: 视盘二值掩码
          - cup_mask: 视杯二值掩码
          - output_path: 结果保存路径，若为 None 则不保存
          - show_overlay: 是否显示叠加掩码的结果
          - show_bbox: 是否显示绘制边框的结果
          - alpha: 叠加掩码时的透明度
        """
        results = []

        if show_overlay:
            # 叠加掩码的结果
            disc_overlay = self.overlay_mask(original, disc_mask, color=(0, 255, 0), alpha=alpha)
            cup_overlay = self.overlay_mask(original, cup_mask, color=(0, 255, 0), alpha=alpha)
            results.extend([disc_overlay, cup_overlay])

        if show_bbox:
            # 绘制边框的结果
            disc_bbox = self.draw_bounding_boxes(original, disc_mask, thickness=2)
            cup_bbox = self.draw_bounding_boxes(original, cup_mask, thickness=2)
            results.extend([disc_bbox, cup_bbox])

        # 如果提供了输出路径，则保存结果
        if output_path:
            for i, result in enumerate(results):
                cv2.imencode('.jpg', result)[1].tofile(f"{output_path}_result_{i + 1}.jpg")

        return results


if __name__ == "__main__":
    processor = EyeFundusProcessor()

    # 示例：直接读取图像数组
    img_path = r'/眼底视盘视杯渗出物突出显示/示例.jpg'
    image = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)

    image_clahe, disc_mask, cup_mask, hist_green, hist_red = processor.process_image_from_array(image)

    # 显示或保存处理结果
    results = processor.display_results(image_clahe, disc_mask, cup_mask, output_path='output_image', show_overlay=True,
                                        show_bbox=True)